package io.openmessaging.utils;

import java.util.HashSet;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 统计耗时
 * 多线程会取平均数
 *
 * @author cuijd    2021/9/16 14:54
 */
public class Cost {
    private AtomicLong cost = new AtomicLong();

    private ThreadLocal<Long> startTime = new ThreadLocal<>();
    HashSet<Thread> threads = new HashSet<>();

    public Cost start() {
        threads.add(Thread.currentThread());
        Assert.isTrue(startTime.get() == null, "禁止连续调用start");
        startTime.set(System.nanoTime());
        return this;
    }

    public void pause() {
//        System.out.println("pause " + System.nanoTime() + " " + startTime.get() + " "+(System.nanoTime() - startTime.get()));
        cost.addAndGet(System.nanoTime() - startTime.get());
        startTime.set(null);
    }

    /**
     * 返回累计的耗时
     *
     * @return 毫秒数
     */
    public long cost() {
        return threads.size() == 0 ? 0 : cost.get() / 1000 / 1000 / threads.size();
    }

    public interface Function {
        void exec() throws Exception;
    }

    public interface Function2<E> {
        E exec() throws Exception;
    }

    public Cost cost(Function f) {
        try {
            start();
            f.exec();
        } catch (Exception e) {
            if (e instanceof RuntimeException)
                throw (RuntimeException) e;
            else throw new RuntimeException(e);
        } finally {
            pause();
        }
        return this;
    }

    public <E> E cost(Function2<E> f) {
        try {
            start();
            return f.exec();
        } catch (Exception e) {
            throw new RuntimeException(e);
        } finally {
            pause();
        }
    }

    @Override
    public String toString() {
        return Long.toString(cost()) + " ms ";
    }
}
